from ucimlrepo import fetch_ucirepo
import pandas as pd
from itertools import combinations

# 获取数据集
wine_quality = fetch_ucirepo(id=186)

# 数据 (作为pandas dataframes)
X = wine_quality.data.features
y = wine_quality.data.targets

# 将数据转换为适合Apriori算法处理的格式
transactions = X.applymap(lambda x: 1 if x > 0 else 0).values.tolist()


def apriori(transactions, min_support):
    """
    实现Apriori算法生成频繁项集。

    参数:
    transactions -- 事务列表，每个事务是一个包含项的集合
    min_support -- 最小支持度阈值

    返回:
    frequent_itemsets -- 频繁项集及其支持度列表
    """
    # 第一步：生成候选项集
    items = set()
    for transaction in transactions:
        for item in transaction:
            items.add(frozenset([item]))

    # 第二步：计算支持度并筛选频繁项集
    frequent_itemsets = []
    current_itemsets = list(items)
    k = 1

    while current_itemsets:
        itemset_support = {}
        for transaction in transactions:
            for itemset in current_itemsets:
                if itemset.issubset(transaction):
                    if itemset not in itemset_support:
                        itemset_support[itemset] = 0
                    itemset_support[itemset] += 1

        current_frequent_itemsets = []
        for itemset, support in itemset_support.items():
            support_count = support / len(transactions)
            if support_count >= min_support:
                current_frequent_itemsets.append(itemset)
                frequent_itemsets.append((itemset, support_count))

        # 生成新的候选项集
        current_itemsets = []
        for i in range(len(current_frequent_itemsets)):
            for j in range(i + 1, len(current_frequent_itemsets)):
                union_set = current_frequent_itemsets[i].union(current_frequent_itemsets[j])
                if len(union_set) == k + 1:
                    current_itemsets.append(union_set)
        k += 1

    return frequent_itemsets


def generate_rules(frequent_itemsets, min_confidence):
    """
    从频繁项集中生成关联规则。

    参数:
    frequent_itemsets -- 频繁项集及其支持度列表
    min_confidence -- 最小置信度阈值

    返回:
    rules -- 关联规则列表，每条规则由前件、后件和置信度组成
    """
    rules = []
    for itemset, support in frequent_itemsets:
        for i in range(1, len(itemset)):
            for antecedent in combinations(itemset, i):
                antecedent = frozenset(antecedent)
                consequent = itemset - antecedent
                confidence = support / sum(sup for itm, sup in frequent_itemsets if itm == antecedent)
                if confidence >= min_confidence:
                    rules.append((antecedent, consequent, confidence))
    return rules


# 设置参数
min_support = 0.01
min_confidence = 0.5

# 运行Apriori算法
frequent_itemsets = apriori(transactions, min_support)
rules = generate_rules(frequent_itemsets, min_confidence)

# 输出结果（中文）
print("频繁项集：")
for itemset, support in frequent_itemsets:
    print(f"{set(itemset)}: {support}")

print("\n关联规则：")
for antecedent, consequent, confidence in rules:
    print(f"{set(antecedent)} => {set(consequent)} (置信度: {confidence})")
